// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <turbo/flags/flag.h>

#include <melon/executors/io_thread_pool_executor.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/common/memory/mmap_allocator.h>
#include <pollux/exec/spiller.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/type/type.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_DECLARE_FLAG(std::string, spiller_benchmark_compression_kind);
TURBO_DECLARE_FLAG(std::string, spiller_benchmark_name);
TURBO_DECLARE_FLAG(std::string, spiller_benchmark_path);
TURBO_DECLARE_FLAG(std::string, spiller_benchmark_spiller_type);
DECLARE_uint32(spiller_benchmark_num_key_columns);
DECLARE_uint32(spiller_benchmark_num_spill_vectors);
DECLARE_uint32(spiller_benchmark_spill_executor_size);
DECLARE_uint32(spiller_benchmark_spill_vector_size);
DECLARE_uint64(spiller_benchmark_max_spill_file_size);
DECLARE_uint64(spiller_benchmark_min_spill_run_size);
DECLARE_uint64(spiller_benchmark_write_buffer_size);

namespace kumo::pollux::exec::test {
/// This test measures the spill input overhead in spill join & probe.
class SpillerBenchmarkBase {
 public:
  SpillerBenchmarkBase() = default;

  virtual ~SpillerBenchmarkBase() = default;

  /// Sets up the test.
  virtual void setUp() = 0;

  /// Runs the test.
  virtual void run() = 0;

  /// Prints out the measured test stats.
  virtual void printStats() const;

  /// Cleans up the test.
  virtual void cleanup();

 protected:
  std::shared_ptr<pollux::memory::MemoryPool> rootPool_;
  std::shared_ptr<pollux::memory::MemoryPool> pool_;
  RowTypePtr rowType_;
  uint32_t numInputVectors_;
  uint32_t inputVectorSize_;
  std::unique_ptr<VectorFuzzer> vectorFuzzer_;
  std::vector<RowVectorPtr> rowVectors_;
  std::unique_ptr<melon::IOThreadPoolExecutor> executor_;
  std::shared_ptr<exec::test::TempDirectoryPath> tempDir_;
  std::string spillDir_;
  std::shared_ptr<filesystems::FileSystem> fs_;
  std::unique_ptr<SpillerBase> spiller_;
  // Stats.
  uint64_t executionTimeUs_{0};
  melon::Synchronized<common::SpillStats> spillStats_;
};
} // namespace kumo::pollux::exec::test
